package com.liu.ai.chat;

import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.http.client.reactive.ClientHttpRequest;
import org.springframework.web.client.RestClient;
import org.springframework.web.reactive.function.BodyInserter;
import org.springframework.web.reactive.function.BodyInserters;
import org.springframework.web.reactive.function.client.ExchangeFilterFunction;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.util.function.BiFunction;
import java.util.function.Consumer;
import java.util.function.Function;

@Slf4j
public class Chat {


    public static void main(String[] args) throws InterruptedException {


        OllamaApi ollamaApi = new OllamaApi("http://localhost:11434",
                RestClient.builder(),
                WebClient.builder()
                    .filter(logRequest())
                    .filter(logResponse())
        );
        OllamaOptions options = OllamaOptions.builder()
                .model("qwen2.5:3b")
                .temperature(0.9)
                .build();

        OllamaChatModel chatModel = OllamaChatModel.builder()
                .ollamaApi(ollamaApi)
                .defaultOptions(options)
                .build();
        Flux<ChatResponse> res = chatModel.stream(new Prompt("介绍一下你自己"));


        StringBuilder sb = new StringBuilder();

        res.doOnNext(i-> sb.append(i.getResult().getOutput().getText()))
            .reduce("",new BiFunction< String,ChatResponse, String>() {
                @Override
                public String apply(String s, ChatResponse chatResponse) {
        //                System.out.println(Thread.currentThread().getName() + "  111 "+ s);
                    return s + chatResponse.getResult().getOutput().getText();
                }
            })
            .subscribe(System.out::println, Throwable::printStackTrace, new Runnable() {
                @Override
                public void run() {

                    System.out.println("====================消费完成执行============== \n" + sb.toString());
                }
            });


        // 消费两次，请求了两次
//        res.subscribe(item -> {
//            item.getResults().forEach(i -> System.out.println(Thread.currentThread().getName() + "  222 " + i.getOutput().getContent()));
//        });
        Thread.sleep(1000 * 60 * 40);
    }

    public static ExchangeFilterFunction logRequest() {
        return ExchangeFilterFunction.ofRequestProcessor(clientRequest -> {
            System.out.println("Request: " + clientRequest.method() + " " + clientRequest.url());


            clientRequest.headers().forEach((name, values) ->
                    values.forEach(value -> System.out.println(name + ": " + value)));

            BodyInserter<?, ? super ClientHttpRequest> body = clientRequest.body();
            if (clientRequest.body() instanceof BodyInserter<?, ?>) {
                // 根据不同类型处理
                if (body instanceof BodyInserters.FormInserter) {
                    BodyInserters.FormInserter<?> formInserter = (BodyInserters.FormInserter<?>) body;
                    System.out.println("Form data: " + formInserter);
                } else if (body instanceof BodyInserters.MultipartInserter) {
                    BodyInserters.MultipartInserter multipartInserter = (BodyInserters.MultipartInserter) body;
                    System.out.println("Multipart data: " + multipartInserter);
                }
            }

            System.out.println(clientRequest.body());
            return Mono.just(clientRequest);
        });
    }

    public static ExchangeFilterFunction logResponse() {
        return ExchangeFilterFunction.ofResponseProcessor(clientResponse -> {
            System.out.println("Response status: " + clientResponse.statusCode());
            clientResponse.headers().asHttpHeaders().forEach((name, values) ->
                    values.forEach(value -> System.out.println(name + ": " + value)));
            return Mono.just(clientResponse);
        });
    }
}

